#!/usr/bin/env python
# -*- encoding: UTF-8 -*-
from __future__ import print_function

from scipy import weave
import numpy as np
import cv2
import math

def _thinningIteration(im, iter):
    I, M = im, np.zeros(im.shape, np.uint8)
    expr = """
    for (int i = 1; i < NI[0]-1; i++) {
        for (int j = 1; j < NI[1]-1; j++) {
            int p2 = I2(i-1, j);
            int p3 = I2(i-1, j+1);
            int p4 = I2(i, j+1);
            int p5 = I2(i+1, j+1);
            int p6 = I2(i+1, j);
            int p7 = I2(i+1, j-1);
            int p8 = I2(i, j-1);
            int p9 = I2(i-1, j-1);
            int A  = (p2 == 0 && p3 == 1) + (p3 == 0 && p4 == 1) +
                     (p4 == 0 && p5 == 1) + (p5 == 0 && p6 == 1) +
                     (p6 == 0 && p7 == 1) + (p7 == 0 && p8 == 1) +
                     (p8 == 0 && p9 == 1) + (p9 == 0 && p2 == 1);
            int B  = p2 + p3 + p4 + p5 + p6 + p7 + p8 + p9;
            int m1 = iter == 0 ? (p2 * p4 * p6) : (p2 * p4 * p8);
            int m2 = iter == 0 ? (p4 * p6 * p8) : (p2 * p6 * p8);
            if (A == 1 && B >= 2 && B <= 6 && m1 == 0 && m2 == 0) {
                M2(i,j) = 1;
            }
        }
    }
    """

    weave.inline(expr, ["I", "iter", "M"])
    return (I & ~M)


def thinning(src):
    dst = src.copy() / 255
    prev = np.zeros(src.shape[:2], np.uint8)
    diff = None

    while True:
        dst = _thinningIteration(dst, 0)
        dst = _thinningIteration(dst, 1)
        diff = np.absolute(dst - prev)
        prev = dst.copy()
        if np.sum(diff) == 0:
            break

    return dst * 255

def houghLine(src):
    cv2.imshow("source", src)
    dst = cv2.cvtColor(src, cv2.COLOR_BGR2GRAY) # 灰度图
    dst = cv2.pyrDown(dst) # 降阶一次
    dst = cv2.pyrDown(dst) # 降阶两次
    retval,dst = cv2.threshold(dst,10,255,cv2.THRESH_BINARY_INV) # 二值化
    dst = cv2.morphologyEx(dst, cv2.MORPH_OPEN, np.ones((5,5),np.uint8)) # 开运算
    dst = cv2.morphologyEx(dst, cv2.MORPH_CLOSE, np.ones((5,5),np.uint8)) # 闭运算
    cv2.imshow("dst", dst)
    dst = thinning(dst)
    cdst = cv2.cvtColor(dst, cv2.COLOR_GRAY2BGR)

    if True: # HoughLinesP
        lines = cv2.HoughLinesP(dst, 1, math.pi/180.0, 50, np.array([]), 20, 10)
        if lines is None:return None
        a,b,c = lines.shape
        for i in range(a):
            for j in range(b):
                cv2.line(cdst, (lines[i][j][0], lines[i][j][1]), (lines[i][j][2], lines[i][j][3]), (0, 0, 255), 3)
    else:    # HoughLines
        lines = cv2.HoughLines(dst, 1, math.pi/180.0, 50, np.array([]), 0, 0)
        if lines is None:return None
        if lines is not None:
            a,b,c = lines.shape
            for i in range(a):
                rho = lines[i][0][0]
                theta = lines[i][0][1]
                a = math.cos(theta)
                b = math.sin(theta)
                x0, y0 = a*rho, b*rho
                pt1 = ( int(x0+1000*(-b)), int(y0+1000*(a)) )
                pt2 = ( int(x0-1000*(-b)), int(y0-1000*(a)) )
                cv2.line(cdst, pt1, pt2, (0, 0, 255), 3)
    cv2.imshow("detected lines", cdst)


def getLine(src):
    dst = cv2.cvtColor(src, cv2.COLOR_BGR2GRAY) # 灰度图
    dst = cv2.pyrDown(dst) # 降阶一次
    dst = cv2.pyrDown(dst) # 降阶两次
    retval,dst = cv2.threshold(dst,10,255,cv2.THRESH_BINARY_INV) # 二值化
    dst = cv2.morphologyEx(dst, cv2.MORPH_OPEN, np.ones((5,5),np.uint8)) # 开运算
    dst = cv2.morphologyEx(dst, cv2.MORPH_CLOSE, np.ones((5,5),np.uint8)) # 闭运算
    dst = thinning(dst)
    lines = cv2.HoughLinesP(dst, 1, math.pi/180.0, 50, np.array([]), 20, 10)
    if lines is None: return None
    a,b,c = lines.shape
    max_dist = 0
    line = None
    for i in range(a):
        for j in range(b):
            pt1 = lines[i][j][0:2]
            pt2 = lines[i][j][2:4]
            dist = np.linalg.norm(pt1-pt2)
            if dist > max_dist:
                max_dist = dist
                pt1 = lines[i][j][0:2]*2**2
                pt2 = lines[i][j][2:4]*2**2
                pt1 = np.hstack((pt1,1))
                pt2 = np.hstack((pt2,1))
                line = np.cross(pt1,pt2) # 直线的齐次表示
    return line # 直线的齐次表示

def main():
    cap = cv2.VideoCapture(0)
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        show_img = frame.copy()
        houghLine(show_img)
        ch = cv2.waitKey(1) & 0xFF
        if ch == 27:
            break

if __name__ == '__main__':
    main()

